from sqlalchemy import create_engine, Table, Column, String, Integer
import asyncio
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore
from llama_index.core.tools import QueryEngineTool
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, VectorStoreIndex
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    select,
    column,
)
from sqlalchemy import insert, create_engine, String, text, Integer
from sqlalchemy.orm import declarative_base, mapped_column

# 创建内存数据库
engine = create_engine("sqlite:///:memory:")
metadata = MetaData()

# 定义表结构
city_stats = Table(
    "city_stats", metadata,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False)
)
metadata.create_all(engine)

# 插入示例数据
data = [
    {"city_name": "Toronto", "population": 2731571, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13929286, "country": "Japan"},
    {"city_name": "Berlin", "population": 600000, "country": "Germany"},
    {"city_name": "Beijing", "population": 1200000000, "country": "China"}
]
with engine.begin() as conn:
    for row in data:
        conn.execute(city_stats.insert().values(**row))

from llama_index.core import SQLDatabase

# 限定仅使用city_stats表（避免全库扫描）
sql_db = SQLDatabase(engine, include_tables=["city_stats"])

from llama_index.core.query_engine import NLSQLTableQueryEngine

sql_query_engine = NLSQLTableQueryEngine(
    sql_database=sql_db,
    tables=["city_stats"],
)
from llama_index.core.tools import QueryEngineTool





sql_tool = QueryEngineTool.from_defaults(
    query_engine=sql_query_engine,
    description=(
        "Useful for translating a natural language query into a SQL query over"
        " a table containing: city_stats, containing the population/country of"
        " each city"
    ),
)
vector_tools = []

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=([sql_tool]),
)


from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema

# 定义表结构节点（可附加额外说明）
table_schema = SQLTableSchema(
    table_name="city_stats",
    context_str="存储城市人口与国家信息，'population'对应人口，'city_name'对应城市"
)

# 创建节点映射器
node_mapping = SQLTableNodeMapping(sql_db)

# 构建向量索引（将表结构存入向量空间）
obj_index = ObjectIndex.from_objects(
    [table_schema],
    node_mapping=node_mapping,
    index_cls=VectorStoreIndex
)

# 组合检索式查询引擎
from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine

retriever_engine = SQLTableRetrieverQueryEngine(
    sql_db,
    table_retriever=obj_index.as_retriever(similarity_top_k=1)
)

from llama_index.core import Settings
from llama_index.core.callbacks import LlamaDebugHandler
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler,TokenCountingHandler

debug_handler = LlamaDebugHandler(print_trace_on_end=True)

tokenCountingHandler=TokenCountingHandler()

callback_manager = CallbackManager(handlers=[debug_handler,tokenCountingHandler])

Settings.callback_manager = callback_manager

# 模糊查询（引擎自动检索匹配表结构）
response = retriever_engine.query("人口最多的城市是哪个？")
print(response)


